import argparse
import subprocess

import numpy as np
import torch
from mmcv import Config
from mmcv.runner import load_checkpoint

from mmdet.models import build_detector


def parse_args():
    parser = argparse.ArgumentParser(description="Process a checkpoint to be published")
    parser.add_argument("in_file", help="input checkpoint filename")
    parser.add_argument("config", help="model config")
    parser.add_argument("out_file", help="output checkpoint filename")
    args = parser.parse_args()
    return args


def process_checkpoint(in_file, config, out_file):
    checkpoint = torch.load(in_file, map_location="cpu")
    # remove optimizer for smaller file size
    if "optimizer" in checkpoint:
        del checkpoint["optimizer"]
    # if it is necessary to remove some sensitive data in checkpoint['meta'],
    # add the code here.
    torch.save(checkpoint, out_file)
    sha = subprocess.check_output(["sha256sum", out_file]).decode()
    if out_file.endswith(".pth"):
        out_file_name = out_file[:-4]
    else:
        out_file_name = out_file
    final_file = out_file_name + f"-{sha[:8]}.pth"
    subprocess.Popen(["mv", out_file, final_file])


def repvgg_model_convert(cfg, pth, save_path=None):
    cfg.model.backbone.deploy = False
    model = build_detector(
        cfg.model, train_cfg=cfg.get("train_cfg"), test_cfg=cfg.get("test_cfg")
    )
    load_checkpoint(model, pth, map_location="cpu")
    converted_weights = {}
    for name, module in model.named_modules():
        if hasattr(module, "repvgg_convert"):
            kernel, bias = module.repvgg_convert()
            converted_weights[name + ".rbr_reparam.weight"] = kernel
            converted_weights[name + ".rbr_reparam.bias"] = bias
        elif isinstance(module, torch.nn.Linear):
            converted_weights[name + ".weight"] = module.weight.detach().cpu().numpy()
            converted_weights[name + ".bias"] = module.bias.detach().cpu().numpy()
        else:
            print(name, type(module))
    #     del model

    cfg.model.backbone.deploy = True
    deploy_model = build_detector(
        cfg.model, train_cfg=cfg.get("train_cfg"), test_cfg=cfg.get("test_cfg")
    )
    for name, param in deploy_model.named_parameters():
        if "backbone" in name:
            print(
                "deploy param: ", name, param.size(), np.mean(converted_weights[name])
            )
            param.data = torch.from_numpy(converted_weights[name]).float()
    deploy_model.neck = model.neck
    deploy_model.bbox_head = model.bbox_head

    if save_path is not None:
        torch.save(deploy_model.state_dict(), save_path)
        out_file = save_path
        sha = subprocess.check_output(["sha256sum", out_file]).decode()
        if out_file.endswith(".pth"):
            out_file_name = out_file[:-4]
        else:
            out_file_name = out_file
        final_file = out_file_name + f"-{sha[:8]}.pth"
        subprocess.Popen(["mv", out_file, final_file])

    return deploy_model


def main():
    args = parse_args()
    cfg = Config.fromfile(args.config)
    repvgg_model_convert(cfg, args.in_file, save_path=args.out_file)


if __name__ == "__main__":
    main()
